﻿using System;
using System.Collections.Generic;
using QSharp.Scheme.Mathematics;

namespace QSharp.Signal.Classic
{
    public static class Fft
    {
        private class FactorAidedInteger
        {
            public FactorAidedInteger(int n)
            {
                N = n;
                _factors = n.Factor();
                _currIdx = 0;
            }

            public int Decompose()
            {
                if (N == 1)
                {
                    return 1;
                }
                var n1 = _factors[_currIdx++];
                N /= n1;
                return n1;
            }

            public void Restore(int n1)
            {
                N *= n1;
                _currIdx--;
            }


            public int N { get; private set; }

            private readonly IList<int> _factors;
            private int _currIdx;
        }

        private static void Dft(this IList<Complex> input, int start, int skip, IList<Complex> output, double a)
        {
            var ni = input.Count;
            var n = output.Count;
            for (var k = 0; k < n; k++)
            {
                var d = Complex.Zero;
                var j = 0;
                for (var i = start; i < ni; i+=skip, j++)
                {
                    d += input[i]*Complex.Polar(-Math.PI*2*j*k*a);
                }
                output[k] = d;
            }
        }

        public static void Dft(this IList<Complex> input, IList<Complex> output)
        {
            var n = input.Count;
            var a = 1.0/n;
            for (var k = 0; k < n; k++)
            {
                var d = Complex.Zero;
                for (var i = 0; i < n; i++)
                {
                    d += input[i] * Complex.Polar(-Math.PI * 2 * i * k * a);
                }
                output[k] = d;
            }
        }

        /// <summary>
        ///  Decomposes an integer number into two integer numbers one of which being
        ///  the greatest factor and the other thereby being the smallest
        /// </summary>
        /// <param name="n">The integer number to decompose</param>
        /// <param name="n1">The smaller factor</param>
        /// <param name="n2">The greater factor</param>
        public static bool FirstPrimeDecompose(int n, out int n1, out int n2)
        {
            var sqrtN = (int) Math.Sqrt(n);
            for (n1 = 2; n1 <= sqrtN; n1++)
            {
                n2 = n/n1;
                if (n1*n2 == n)
                {
                    return true;
                }
            }
            n1 = 1;
            n2 = n;
            return false;
        }

        private static void GetSubsequenceIndexGenerator(int origStart, int origSkip, int subStart, int subSkip, out int start, out int skip)
        {
            start = origStart + origSkip * subStart;
            skip = origSkip*subSkip;
        }

        private static void FftRecursive(this IList<Complex> input, FactorAidedInteger n, int start, int skip, IList<Complex> output)
        {
            var n0 = n.N;
            var n1 = n.Decompose();
            var n2 = n.N;
            if (n1 > 1 && n2 > 1)
            {
                for (var k = 0; k < n0; k++)
                {
                    output[k] = 0;
                }
                for (var r = 0; r < n1; r++)
                {
                    int newStart, newSkip;
                    GetSubsequenceIndexGenerator(start, skip, r, n1, out newStart, out newSkip);

                    // N2-point FFT
                    var tmp = new Complex[n2];
                    input.FftRecursive(n, newStart, newSkip, tmp);

                    for (var k = 0; k < n0; k++)
                    {
                        var k2 = k%n2;
                        output[k] += tmp[k2]*Complex.Polar(-Math.PI*2*r*k/n0);
                    }
                }
            }
            else
            {
                input.Dft(start, skip, output, 1.0/n0);
            }
            n.Restore(n1);
        }

        public static void FftRecursive(this IList<Complex> input, IList<Complex> output)
        {
            var n = new FactorAidedInteger(input.Count);
            input.FftRecursive(n, 0, 1, output);
        }
    }
}
